package org.jiangwei.cmput696.entitylinking.tacindex;

import java.io.File;
import java.io.IOException;
import java.util.LinkedList;
import java.util.List;

import org.apache.lucene.analysis.Analyzer;
import org.apache.lucene.analysis.standard.StandardAnalyzer;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.Field;
import org.apache.lucene.util.Version;
import org.jiangwei.cmput696.entitylinking.Helper;
import org.jiangwei.cmput696.entitylinking.QueryParsingException;
import org.w3c.dom.Element;
import org.w3c.dom.NodeList;

public class TacIndexConstructor {

	private List<File> xmlQueue;

	private TacDocumentPool pool;

	private TacLuceneIndexer indexer;

	public TacIndexConstructor() {
		xmlQueue = new LinkedList<File>();
		String dirPath =
		 "/largedata1/cmput696/knowledge_base/tac-kbp-knowledge-base-2011/data";
		//String dirPath = "../test_files/tac-kbp-sample";
		File dir = new File(dirPath);
		for (File file : dir.listFiles()) {
			if (file.isFile() && file.getName().endsWith(".xml")) {
				xmlQueue.add(file);
			}
		}
	}

	public int getTotal() {
		return xmlQueue.size();
	}

	private void loadTacXML(File file) throws QueryParsingException {
		NodeList nodeLst = Helper.extractXMLNodeList(file, "entity");
		for (int s = 0; s < nodeLst.getLength(); s++) {
			Element elementNode = (Element) nodeLst.item(s);

			String entityId = elementNode.getAttribute("id");
			String tacName = elementNode.getAttribute("wiki_title");
			String wikiName = elementNode.getAttribute("name");
			String entityType = elementNode.getAttribute("type");
			NodeList facts = elementNode.getElementsByTagName("fact");
			NodeList factLinks = elementNode.getElementsByTagName("link");
			String factStr = TacHelper.createFactJSON(facts)
					+ TacHelper.createFactJSON(factLinks);
			// strip the '<![CDATA[' and ']]>' tags
			String wikiText = Helper.readNodeChild(elementNode, "wiki_text");
			if (wikiText.startsWith("<![CDATA[")) {
				wikiText = wikiText.substring(9);
			}
			if (wikiText.endsWith("]]>")) {
				wikiText = wikiText.substring(0, wikiText.length() - 3);
			}

			Document doc = new Document();

			doc.add(new Field(TacHelper.KEY_ID, entityId, Field.Store.YES,
					Field.Index.NOT_ANALYZED_NO_NORMS));
			doc.add(new Field(TacHelper.KEY_TAC_TITLE, tacName,
					Field.Store.YES, Field.Index.NOT_ANALYZED_NO_NORMS));

			doc.add(new Field(TacHelper.KEY_WIKI_TITLE, wikiName,
					Field.Store.YES, Field.Index.NOT_ANALYZED_NO_NORMS));
			doc.add(new Field(TacHelper.KEY_WIKI_TITLE_FOR_SEARCH, wikiName,
					Field.Store.NO, Field.Index.ANALYZED));

			doc.add(new Field(TacHelper.KEY_TEXT, wikiText, Field.Store.YES,
					Field.Index.NO));
			doc.add(new Field(TacHelper.KEY_TYPE, entityType, Field.Store.YES,
					Field.Index.NO));
			doc.add(new Field(TacHelper.KEY_FACT, factStr, Field.Store.YES,
					Field.Index.NO));

			// if (++printCount <= 10) {
			// System.out.println(TacHelper.KEY_ID + ": " + entityId);
			// System.out.println(TacHelper.KEY_TAC_TITLE + ": " + tacName);
			// System.out.println(TacHelper.KEY_WIKI_TITLE + ": " + wikiName);
			// System.out.println(TacHelper.KEY_FACT + ": " + factStr);
			// System.out.println(TacHelper.KEY_TYPE + ": " + entityType);
			// System.out.println(printCount);
			// }

			pool.push(doc);
		}
		System.gc();
	}

	public void createWithMultiThreads(int threadCount, String indexDir)
			throws QueryParsingException {
		pool = new TacDocumentPool();
		try {
			indexer = new TacLuceneIndexer(threadCount, indexDir, pool);
		} catch (IOException e) {
			throw new QueryParsingException("Failed to create Lucene indexer.",
					e);
		}
		indexer.startWorkers();

		for (File file : xmlQueue) {
			loadTacXML(file);
			System.out.println(file.getName());
		}

		pool.setStopped(true);
		indexer.closeWriter();
	}

	/**
	 * @param args
	 */
	public static void main(String[] args) {
		if (args.length < 1) {
			System.out
					.println("usage: TacIndexConstructor <thread_count> [index dir]");
			System.exit(-1);
		}
//		args = new String[] {
//				"1",
//				"../test_files/tac-kbp-sample/index" };
		int threadCount = Integer.parseInt(args[0]);

		TacIndexConstructor tacCtor = new TacIndexConstructor();
		try {
			tacCtor.createWithMultiThreads(threadCount,
					args.length > 1 ? args[1] : null);
		} catch (QueryParsingException e) {
			e.printStackTrace();
		}
	}
}
